1775F - Laboratory on Pluto - CodeForces Solution

#include<iostream>
#include<cstring>
#include<vector>
using namespace std;
using LL = long long;
const int maxn = 1005, INF = 0x3f3f3f3f;
LL g[maxn], f[5][maxn];
int mod;

void init(){
    g[0] = 1;
    for(int i = 1; i < maxn; i++)
        for(int j = i; j < maxn; j++)
            g[j] = (g[j] + g[j - i]) % mod;
    f[0][0] = 1;
    for(int i = 1; i <= 4; i++){
        for(int j = 0; j < maxn; j++)
            for(int k = 0; j + k < maxn; k++)
                f[i][j + k] = (f[i][j + k] + f[i - 1][j] * g[k]) % mod;
    }
}

int main(){

    cin.tie(0);
    cout.tie(0);
    ios::sync_with_stdio(0);

    int T, op;
    cin >> T >> op;
    if (op == 1){
        while(T--){
            int n;
            cin >> n;
            int ans = INF, h;
            for(int i = 1; i <= n; i++){
                int j = (n + i - 1) / i;
                if (i + j < ans){
                    ans = i + j;
                    h = i;
                }
            }
            vector<string> s(h);
            int w = (n + h - 1) / h;
            for(int i = 0, k = 0; i < h; i++){
                s[i].resize(w);
                for(int j = 0; j < w; j++){
                    if (++k > n) s[i][j] = '.';
                    else s[i][j] = '#';
                }
            }
            cout << h << ' ' << w << '\n';
            for(auto &x : s) cout << x << '\n';
        }
    }
    else{
        cin >> mod;
        init();
        while(T--){
            int n;
            cin >> n;
            int ans = INF;
            vector<int> cand;
            for(int i = 1; i <= n; i++){
                int j = (n + i - 1) / i;
                if (i > j) break;
                if (i + j < ans){
                    ans = i + j;
                    cand = {i};
                }
                else if (i + j == ans) cand.push_back(i);
            }

            LL cnt = 0;
            for(auto h : cand){
                int w = (n + h - 1) / h;
                int delta = w * h - n;
                if (w == h) cnt = (cnt + f[4][delta]) % mod;
                else cnt = (cnt + 2 * f[4][delta]) % mod;
            }
            cout << ans * 2 << ' ' << cnt << '\n';
        }
    }

}